import segmentation_models_pytorch as smp
import timm
import torch
from torch.nn import Module
from torch import nn
import torch.nn.functional as F
from models.ms.global_net import GlobalNet
from models.ms.strip_ts_model import StripTSModel

# timm中用于分类的模型
timm_classify_models = ['mobilenetv2_100',
                        'mobilenetv2_140', 'ghostnet_050', 'ghostnet_100', 'ghostnet_130',
                        'mobilenetv3_small_075', 'mobilenetv3_small_100', 'mobilenetv3_large_100',
                        'resnet10t', 'resnet14t', 'resnet18', 'resnet18d', 'resnet26', 'resnet26d', 'resnet26t',
                        'resnet32ts',
                        'resnet33ts', 'resnet34', 'resnet34d', 'resnet50', 'resnet50_gn', 'resnet50c', 'resnet50d',
                        'resnet50s',
                        'resnet50t', 'resnet51q', 'resnet61q', 'resnet101', 'resnet101c', 'resnet101d', 'resnet101s',
                        'resnet152',
                        'resnet152c', 'resnet152d', 'resnet152s', 'resnet200', 'resnet200d', 'resnetaa34d',
                        'resnetaa50',
                        'resnetaa50d', 'resnetaa101d', 'resnetblur18', 'resnetblur50', 'resnetblur50d',
                        'resnetblur101d',
                        'resnetrs50',
                        'resnetrs101', 'resnetrs152', 'resnetrs200', 'resnetrs270', 'resnetrs350', 'resnetrs420',
                        'resnetv2_50',
                        'resnetv2_50d', 'resnetv2_50d_evos', 'resnetv2_50d_frn', 'resnetv2_50d_gn', 'resnetv2_50t',
                        'resnetv2_50x1_bit', 'resnetv2_50x3_bit', 'resnetv2_101', 'resnetv2_101d', 'resnetv2_101x1_bit',
                        'resnetv2_101x3_bit', 'resnetv2_152', 'resnetv2_152d', 'resnetv2_152x2_bit',
                        'resnetv2_152x4_bit']

smp_segment_models = [
    'resnet34', 'resnet50', 'resnext50_32x4d', 'se_resnet50', 'se_resnext50_32x4d',
    'dpn68', 'efficientnet_b4', 'efficientnet_b3'
]

unsupervised_models = ['strip_ts_Model']


class SegmentModel(Module):
    """分割模型"""

    def __init__(self, model_name, pretrained=False, class_num=4, **kwargs):
        super(SegmentModel, self).__init__()
        self.model_name = model_name
        self.class_num = class_num
        if model_name in smp_segment_models:
            if pretrained:
                assert kwargs.get('encoder_weights') is not None, 'a parameter named `encoder_weights` is missing'
            else:
                kwargs['encoder_weights'] = None
            self.model = smp.Unet(model_name, classes=class_num, **kwargs)
        else:
            raise ValueError(f"Model {model_name} is not recognized.")

    def forward(self, x):
        return self.model(x)


class ClassifyModel(Module):
    """分类模型"""

    def __init__(self, model_name, pretrained=False, class_num=4, **kwargs):
        super(ClassifyModel, self).__init__()

        if model_name in timm_classify_models:
            self.encoder = timm.create_model(model_name, pretrained=pretrained, num_classes=class_num, **kwargs)
        elif model_name == 'global_net':
            self.encoder = GlobalNet(num_classes=class_num, **kwargs)
        else:
            raise ValueError(f"Model {model_name} is not recognized.")

    def forward(self, x):
        return self.encoder(x)


class UnsupervisedModel(nn.Module):
    def __init__(self, model_name, **kwargs):
        super(UnsupervisedModel, self).__init__()
        if model_name == 'strip_ts_model':
            self.unsupervised_model = StripTSModel(**kwargs)

    def forward(self, x):
        return self.unsupervised_model(x)


if __name__ == "__main__":
    x = torch.Tensor(1, 3, 256, 1600)
    y = torch.ones(1, 4)

    # test classify 模型
    class_net = ClassifyModel(model_name='global_net', class_num=4, base_dim=30)
    output = class_net(x)
    print(output.size())

    # test segment 模型
    segment_net = SegmentModel(model_name='resnet34', pretrained=True, class_num=4, activation=None,
                               encoder_weights="imagenet")
    output = segment_net(x)
    print(output.size())
